"""
Copyright 2013 Steven Diamond

This file is part of CVXPY.

CVXPY is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

CVXPY is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with CVXPY.  If not, see <http://www.gnu.org/licenses/>.
"""

from cvxpy import *
import create_graph as g
from max_flow import Edge, Node
import pickle
import random as r
import cvxopt

# Multi-commodity flow.
COMMODITIES = 5 # Number of commodities.
r.seed(1)

class MultiEdge(Edge):
    """ An undirected, capacity limited edge with multiple commodities. """
    def __init__(self, capacity):
        self.capacity = capacity
        self.in_flow = Variable(COMMODITIES)
        self.out_flow = Variable(COMMODITIES)

    # Returns the edge's internal constraints.
    def constraints(self):
        return [self.in_flow + self.out_flow == 0,
                sum(abs(self.in_flow)) <= self.capacity]

class MultiNode(Node):
    """ A node with a target flow accumulation and a capacity. """
    def __init__(self, capacity=0):
        self.capacity = capacity
        self.edge_flows = []

    # The total accumulation of flow.
    def accumulation(self):
        return sum(f for f in self.edge_flows)
    
    def constraints(self):
        return [abs(self.accumulation()) <= self.capacity]


if __name__ == "__main__":
    # Read a graph from a file.
    f = open(g.FILE, 'r')
    data = pickle.load(f)
    f.close()

    # Construct nodes.
    node_count = data[g.NODE_COUNT_KEY]
    nodes = [MultiNode() for i in range(node_count)]
    # Add a source and sink for each commodity.
    sources = []
    sinks = []
    for i in range(COMMODITIES):
        source,sink = r.sample(nodes, 2)
        # Only count accumulation of a single commodity.
        commodity_vec = cvxopt.matrix(0,(COMMODITIES,1))
        commodity_vec[i] = 1
        # Initialize the source.
        source.capacity = commodity_vec*Variable()
        sources.append(source)
        # Initialize the sink.
        sink.capacity = commodity_vec*Variable()
        sinks.append(sink)

    # Construct edges.
    edges = []
    for n1,n2,capacity in data[g.EDGES_KEY]:
        edges.append(MultiEdge(capacity))
        edges[-1].connect(nodes[n1], nodes[n2])

    # Construct the problem.
    objective = Maximize(sum(sum(s.accumulation() for s in sinks)))
    constraints = []
    for o in nodes + edges:
        constraints += o.constraints()
    p = Problem(objective, constraints)
    result = p.solve()
    print "Objective value = %s" % result
    # Show how the flow for each commodity.
    for i,s in enumerate(sinks):
        accumulation = sum(f.value[i] for f in s.edge_flows)
        print "Accumulation of commodity %s = %s" % (i, accumulation)
